# encoding=utf-8

import torch
from PIL import Image
from torchvision import transforms
import numpy as np
import cv2 as cv
import os
import functools
from torch import nn
from multiprocessing.dummy import Pool as ThreadPool

from seg_common import PredictFactory
from seg_common import CommonAnnotation
from seg_system import ApplicationConfig
from seg_common.annotation import TimeRecorder
from seg_system.enhance.service.stillgan.models.networks import ResnetGenerator


def tensor2im(input_image, imtype=np.uint8):
    """"Converts a Tensor array into a numpy image array.

    Parameters:
        input_image (tensor) --  the input image tensor array
        imtype (type)        --  the desired type of the converted numpy array
    """
    image_numpy = input_image
    if image_numpy.shape[0] == 1:  # grayscale to RGB
        image_numpy = np.tile(image_numpy, (3, 1, 1))
    image_numpy = (np.transpose(image_numpy, (1, 2, 0)) + 1) / 2.0 * 255.0  # post-processing: tranpose and scaling
    return image_numpy.astype(imtype)


class Identity(nn.Module):
    def forward(self, x):
        return x


def get_norm_layer(norm_type='instance'):
    """Return a normalization layer

    Parameters:
        norm_type (str) -- the name of the normalization layer: batch | instance | none

    For BatchNorm, we use learnable affine parameters and track running statistics (mean/stddev).
    For InstanceNorm, we do not use learnable affine parameters. We do not track running statistics.
    """
    if norm_type == 'batch':
        norm_layer = functools.partial(
            nn.BatchNorm2d, affine=True, track_running_stats=True)
    elif norm_type == 'instance':
        norm_layer = functools.partial(
            nn.InstanceNorm2d, affine=False, track_running_stats=False)
    elif norm_type == 'none':
        def norm_layer(x):
            return Identity()
    else:
        raise NotImplementedError(
            'normalization layer [%s] is not found' % norm_type)
    return norm_layer


class StillganService(PredictFactory):
    # worker = ThreadPool(ApplicationConfig.SystemConfig.NUM_WORKER)

    @CommonAnnotation.override()
    def before(self, obj, **kwargs):
        mode = kwargs.get('mode', ApplicationConfig.SystemConfig.MODE)

        if mode == ApplicationConfig.SystemConfig.MODE_FOLDER:
            return obj
        elif mode == ApplicationConfig.SystemConfig.MODE_PICTURE:
            image = Image.open(obj)
            gray = image.convert('L')
            trans = self.transform(gray)
            trans = trans.unsqueeze(0)
            trans = trans.float()
            trans = trans.to(ApplicationConfig.SystemConfig.DEVICE)
            return trans

    # @TimeRecorder.time_record()
    @CommonAnnotation.override()
    def forward(self, obj, **kwargs):
        with torch.no_grad():
            output = self.model(obj)
            return output

    # @TimeRecorder.time_record()
    @CommonAnnotation.override()
    def after(self, obj, **kwargs):
        o = obj.detach().cpu().numpy()
        l = []

        for each_o in o:
            each_o = tensor2im(each_o)
            (r, g, b) = cv.split(each_o)
            image = cv.merge([b, g, r])
            resize = cv.resize(image, (384, 384))
            l.append(resize)

        return l

    def to_list(self, obj):
        if not isinstance(obj, list):
            return [obj]
        return obj

    def save(self, obj, path, name, **kwargs):
        obj = self.to_list(obj)
        name = self.to_list(name)

        for each_o, each_n in zip(obj, name):
            file_name = os.path.join(path, each_n)
            cv.imwrite(file_name, each_o)
            assert os.path.exists(file_name) is True

    def __init__(self, model_path: str):
        model_d = torch.load(model_path)
        self.model = ResnetGenerator(1, 1, 64, norm_layer=get_norm_layer(),
                                     use_dropout=False, n_blocks=9)
        self.model.load_state_dict(model_d)
        self.model.to(ApplicationConfig.SystemConfig.DEVICE)
        self.model.eval()
        self.transform = transforms.Compose([
            transforms.Resize([518, 518]),
            transforms.ToTensor(),
            transforms.Normalize((0.5), (0.5))
        ])

    def __new__(cls, *args, **kwargs):
        """实际上就是一个标准流程工具类，所以声明成一个单例
        """
        if not hasattr(cls, '_inst'):
            cls._inst = super(StillganService, cls).__new__(cls)
        return cls._inst
